package com.zhang.antlr4demo.hive;

import com.zhang.antlr4demo.bean.SourceTarget;
import org.antlr.v4.runtime.CharStreams;
import org.antlr.v4.runtime.CodePointCharStream;
import org.antlr.v4.runtime.CommonTokenStream;

import java.util.ArrayList;
import java.util.List;

public class TableLineage extends SQLBaseVisitor<Object> {

    private final List<String> inputTables = new ArrayList<>();

    private String outputTable;

    @Override
    public Object visitInsertStatement(SQLParser.InsertStatementContext ctx) {
        outputTable = ctx.T_QIDENTIFIER().getText();
        return super.visitInsertStatement(ctx);
    }

    @Override
    public Object visitFromClause(SQLParser.FromClauseContext ctx) {
        inputTables.add(ctx.T_QIDENTIFIER().getText());
        return super.visitFromClause(ctx);
    }

    @Override
    public Object visitJoinClause(SQLParser.JoinClauseContext ctx) {
        inputTables.add(ctx.T_QIDENTIFIER().getText());
        return super.visitJoinClause(ctx);
    }

    public static List<SourceTarget> lineage(String sql) {
        CodePointCharStream charStream = CharStreams.fromString(sql);
        SQLLexer lexer = new SQLLexer(charStream);
        CommonTokenStream tokens = new CommonTokenStream(lexer);
        SQLParser parser = new SQLParser(tokens);
        SQLParser.StatementContext tree = parser.statement();
        TableLineage lineage = new TableLineage();
        lineage.visit(tree);

        List<SourceTarget> edges = new ArrayList<>();
        for (String source : lineage.inputTables) {
            SourceTarget oneEdge = new SourceTarget();
            oneEdge.setInputTable(source);
            oneEdge.setOutputTable(lineage.outputTable);
            edges.add(oneEdge);
        }
        System.out.println(edges);
        return edges;

    }

    public static void main(String[] args) {
        String sql = "INSERT INTO gmall.tableA SELECT id FROM gmall.tableTwo JOIN gmall.tableThree";
        lineage(sql);
    }
}
